import asyncio
import websockets
import io
from PIL import Image


class Receiver:
    def __init__(self, host: str, port: int):
        self.host    = host
        self.port    = port

    async def get_picture(self, num_of_frames: int):
        images = []
        total  = 0

        async def handler(websocket: websockets.WebSocketClientProtocol):
            nonlocal total
            await websocket.send(f"{num_of_frames}")
            for i in range(num_of_frames):
                image_bytes   = await websocket.recv()
                image         = Image.open(io.BytesIO(image_bytes))
                images.append(image)
                total += 1
                print(total, num_of_frames)

        server = websockets.serve(handler, self.host, self.port)
        await server
        while total < num_of_frames:
            await asyncio.sleep(0.1)
        print('Frames received')
        print('Saving...')
        for i, frame in enumerate(images):
            frame.save(f'./frames/{i}.png')
        print('Complete')

    def __str__(self) -> str:
        return f"Receiver: {self.host}:{self.port}"


class ReceiverManager:
    def __init__(self, ports: list[int], host: str = "localhost"):
        self.ports     = ports
        self.host      = host
        self.receivers = [
            Receiver(
                host = self.host,
                port = port,
            ) for port in self.ports
        ]

    async def get_pictures(self, num_of_frames: int):
        coroutines = [receiver.get_picture(num_of_frames) for receiver in self.receivers]
        await asyncio.gather(*coroutines)

    def __len__(self) -> int:
        return len(self.ports)

    def __str__(self) -> str:
        return f"ReceiverManager: {self.__len__()} ports"

    @staticmethod
    def expand_n_ports(base: int = 8500, n: int = 1):
        return [base + i for i in range(n)]


__all__ = [
    "Receiver",
    "ReceiverManager",
]